import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# 设置图形和坐标轴
fig, ax = plt.subplots(figsize=(9, 9), facecolor='black')
ax.set_facecolor('black')
ax.set_xlim(0, 400)
ax.set_ylim(0, 400)
ax.set_aspect('equal')
ax.axis('off')  # 关闭坐标轴刻度和标签

# 预计算一些常量和数据 (在循环外计算不变的部分)
i = np.arange(int(4e4))  # 使用 40000 个点
x = np.mod(i, 200).astype(float)
y = i / 200.0
k = x / 8.0 - 12.5
e = y / 8.0 - 12.5
o = (k**2 + e**2) / 169.0
d_base = 0.5 + 5.0 * np.cos(o)

# 初始化散点对象
# 初始颜色数据可以是任意的，因为 update 函数会不断更新
scat = ax.scatter([], [], s=1, c='white', alpha=0.4, marker='o', edgecolors='none')

# 动画更新函数
def update(frame):
    t = frame * np.pi / 30  # 计算当前时间 t

    # 重新计算随时间变化的部分
    d = d_base # d_base 已经计算好，d 在这个动画中不随 t 变化，所以直接用 d_base
    sin_term = np.sin(d * 2 + o + t)
    cos_term = np.cos(d * 3 + o * 9 + t)

    # 计算新的 X 和 Y 数据
    new_x = x + d * k * sin_term + e * np.cos(e + t) + 100
    new_y = y / 4.0 - o * 135 + d * 6.0 * cos_term + 275

    # 计算颜色数据 CData
    color_intensity = (d * np.sin(k) * np.sin(t * 4 + e))**2
    # CData 在 matplotlib 中通常需要是 (N, 3) 或 (N, 4) 形状，或者使用单一颜色数组
    # 这里简化为灰度强度，映射到 0-1 的 alpha 值或颜色值
    # 我们使用原始的强度值作为颜色的某种表现，或者只更新 alpha
    # 为了匹配原MATLAB效果，我们尝试直接设置颜色
    # 注意：MATLAB的CData是(N,3)矩阵，这里需要处理
    # 方案1: 计算RGB颜色
    # color_values = np.column_stack([color_intensity, color_intensity, color_intensity])
    # 方案2: 继续使用单一颜色，但可以基于强度调整alpha
    # scat.set_offsets(np.column_stack((new_x, new_y)))
    # scat.set_alpha(...) # 不能直接对每个点设置不同alpha
    # 方案3: 使用colormap和scat.set_array
    # scat.set_array(color_intensity) # 需要预先设置cmap
    # scat.set_offsets(np.column_stack((new_x, new_y)))
    # 为了更接近原MATLAB效果，我们使用 set_array 和一个颜色映射
    # 但为了简单起见，先用基础颜色和强度作为颜色值的一部分
    # 实际上，set_array 是最接近MATLAB CData的方式
    # 需要在初始化时指定cmap
    # 重新初始化scat以使用cmap
    # scat = ax.scatter([], [], s=1, c='white', cmap='gray', alpha=0.4, marker='o', edgecolors='none')
    # 更好的方式是在 update 中使用 set_array，但需要在初始化时指定 c
    # 初始化时指定一个临时颜色和cmap
    # scat = ax.scatter([], [], s=1, c=color_intensity, cmap='gray', alpha=0.4, marker='o', edgecolors='none', vmin=0, vmax=1)
    # 但是初始为空数组时，c不能是空的
    # 我们先用一个临时值初始化，然后在update中设置
    # 为了简化，我们不使用CData的颜色变化，或者用一个近似方法
    # 最接近原MATLAB CData的方式是使用 set_array
    # 我们重新组织代码，初始化时先不设置c，然后在update中设置
    # 但是，如果初始时没有颜色，set_array可能无法工作
    # 最终决定：使用 set_array 和预设的cmap来模拟CData
    if not hasattr(update, 'scat_initialized'):
        # 第一次调用时，设置初始数据
        temp_color = np.zeros_like(new_x) # 临时颜色数据
        scat.set_offsets(np.column_stack((new_x, new_y)))
        scat.set_array(temp_color)
        scat.set_clim(0, 1) # 设置颜色映射的范围
        update.scat_initialized = True
    else:
        scat.set_offsets(np.column_stack((new_x, new_y)))
        scat.set_array(color_intensity)

    return scat,

# 重新初始化scat以配合set_array
scat = ax.scatter([], [], s=1, marker='o', edgecolors='none', alpha=0.4, cmap='gray', vmin=0, vmax=1)

# 创建动画
ani = FuncAnimation(fig, update, interval=20, blit=False, cache_frame_data=False) # interval 20ms 大约 50 FPS

# 显示图形窗口
plt.tight_layout()
plt.show()

# 注意：在某些环境中，plt.show() 会阻塞程序，动画会在此处运行
# 在交互式环境（如 Jupyter Notebook）中，可能需要特殊的后端支持




